package edu.fsu.cs.buzbee.steggo;

import java.awt.Point;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;

import javax.imageio.ImageIO;

public class Encoder {
	// Size of preamble (meta data) in bytes.
	// Preamble:
	//			4 - Magic Number
	//         	4 - Data file size
	public static final int MAGIC = 0xF8E137AA;
	public static final int PREAMBLE_SIZE = 8;
	
	BufferedImage image;
	boolean useAlpha = false;
	
	int bits = 0; // Number of LSBs to use in encoding. If 0, calculated automatically
	/**
	 * Creates a new Encoder capable of encoding data files into a base image
	 * @param base The base image into which data will be encoded
	 */
	public Encoder(BufferedImage base)
	{
		
		// Argument sanity check before accepting construction
		if (base == null)
			throw new NullPointerException("Base image cannot be null");
		if (base.getHeight() == 0 || base.getWidth() == 0)
			throw new RuntimeException("Base image has a dimension value 0");
		
		image = base;
		
	}
	
	/**
	 * Sets the number of LSBs of the 8-bit RGB component used to encode.
	 * bitc must be between 1 and 8. The smaller the number, the less the distortion, but the less data that can be encoded.
	 * @param bitc Number of bits per RGB component to encode data into
	 * @see Encoder.getEncodeBits
	 */
	public void setEncodeBits(int bitc)
	{
		if (bitc < 1 || bitc > 8)
			throw new IndexOutOfBoundsException("Bit count must be between 1 and 8");
		bits = bitc;
	}
	
	/**
	 * Returns the number of LSBs of RGB values used to encode. 
	 * @see Encoder.setEncodeBits 
	 * @return
	 */
	public int getEncodeBits()
	{
		return bits;
	}
	
	/**
	 * Gets the maximum number of bytes that can be encoded in this image using the "bits" LSB
	 * @param bits Number of LSB of RGB to use.  Must be between 1 and 8
	 * @return Number of bytes that can be encoded using this many bits per RGB value
	 */
	int getMaxSize(int bits)
	{
		if (bits > 8 || bits < 1)
			return 0;
		
		return ((image.getHeight() * image.getWidth() * (useAlpha ? 4 : 3) * bits ) / 8) - (PREAMBLE_SIZE);
	}
	
	/**
	 * Encodes the given data file into the base image and writes the result to a new file.
	 * Does not override the base image on the disk, however it will modified the BufferedImage in memory.
	 * Calculates getEncodedBits() if value is out of range.  If this value is 8, the image will not remain!
	 * @param The data file to encode
	 * @param The new encoded image file
	 * @return The encoded version of the base image
	 */
	public void encode(File f, File destination) throws IOException
	{
		if (f == null)
			throw new NullPointerException("Cannot encode null data file");
		if (f.isDirectory())
			throw new IOException("Cannot encode a directory");
		if (!f.exists())
			throw new IOException("Data file not found");
		if (f.length() == 0)
			throw new IOException("Data file is empty");
		if (!f.canRead())
			throw new IOException("Cannot read data file - verify permissions");
		if (!destination.exists() && !destination.createNewFile())
			throw new IOException("Cannot create destination file.");
		if (!destination.canWrite())
			throw new IOException("Cannot write to destination file - verify permissions");
		
		// If the number of LSBs to use is out of range of too small, calculate the minimum needed to store the full data file
		if (bits > 8 || bits < 1 || getMaxSize(bits) < f.length())
		{
			for (int i = 1; i < 9; ++i)
				if (getMaxSize(i) >= f.length())
				{
					bits = i; break;
				}
		}
		
		
		// If the file will not fit in the image (presumably, bits = 8 or 0) return error.
		if (getMaxSize(bits) < f.length())
			throw new IOException("The base image is not large enough to encode this file.");
		if ((image.getWidth() * image.getHeight()) < f.length() + PREAMBLE_SIZE)
			throw new IOException("The encoding algorithm requires at minimum 1 pixel per byte - the data file will not fit in the base image.");
		
		
		//encodeByte(457,1);
		
	//	System.out.printf("First 3 RGB values are: %X %X %X\n", image.getRGB(0, 0), image.getRGB(0, 1), image.getRGB(0, 2) );
		
		long offset = 0;
		offset = encodeInt(offset, MAGIC);
		offset = encodeInt(offset, (int)f.length());
		
		InputStream input = new FileInputStream(f);
		int b;
		while ((b = input.read()) != -1)
			offset = encodeByte(offset, b);
		input.close();
		
		
		ImageIO.write(image, "png", destination);
		
		//System.out.printf("First 3 RGB values after encoding are: %X %X %X\n", image.getRGB(0, 0), image.getRGB(1, 0), image.getRGB(2, 0) );
		

	}
	
	long encodeInt(long offset, int data)
	{
		offset = encodeByte(offset, data & 0xFF);
		offset = encodeByte(offset, (data >>> 8) & 0xFF);
		offset = encodeByte(offset, (data >>> 16) & 0xFF);
		offset = encodeByte(offset, (data >>> 24) & 0xFF);
		return offset;
	}
	
	/**
	 * Encodes byte into base image starting at "offset".  Offset is the pixel number in the file is laid out linearly from 0 to WIDTH*HEIGHT - 1
	 * @param offset Offset to encode byte
	 * @return Offset after encoding.  Should be offset + ceil(8.0/ bits * (3 for RGB or 4 for RGBA))
	 */
	long encodeByte(long offset, int data)
	{
		
		// Sanity checking
		if (data > 0xFF)
			throw new RuntimeException("Cannot encode more than one byte of data using encodeByte");
		if (offset >= image.getWidth() * image.getHeight())
			throw new IndexOutOfBoundsException(offset+" greater than image size: "+image.getWidth() * image.getHeight()+"px");
		if (offset < 0)
			throw new IndexOutOfBoundsException("Offset "+offset+" less than 0");
		 int pixels_to_encode = (int) Math.ceil((8.0/(bits*3)));
		// System.out.println("Pixels to encode: "+pixels_to_encode);
		if ((offset + pixels_to_encode) >= image.getWidth() * image.getHeight()) // Need to see how many pixels it takes to encode a byte and see if we have enough left
			throw new IndexOutOfBoundsException("Not enough bits left in file to encode the byte at offset "+offset + " when encoding in "+bits + " LSBs");
		//System.out.printf("Encoding byte: %X\n", data);
		// Tracks number of bits which have already been encoded as we move through the pixels and components
		int bits_to_encode = 8;
		
		while (bits_to_encode > 0)
		{
			Point p = linearToPoint(offset++);
			// ARGB encoded into the 4 bytes of an int.
			int rgb = image.getRGB(p.x, p.y);
			
		//	System.out.println("Encoding value "+data+" into "+Integer.toString(rgb,16));
			
			int r = rgb >> 16; // Get the R component of the rgb
			r &= 0xFF; // Look at only the least significant byte
			int data_r = (data << (32 - bits)) >>> (32 - bits); // Truncate all but the least significant 'bits' bits of the data
		//	System.out.println("Wants to encode "+data_r + " into "+r);
			r >>>= (bits); r <<= bits; // Set lower 'bits' bits of value to 0 so we can do an or
			r |= data_r; // Encode data_r bits.  Hopefully, only "bits" bits exist.
		//	System.out.println("Comes out with "+r);
			data >>>= bits; // Shift the data so the next least significant 'bits' 'bits' are the next to be encoded.
			bits_to_encode -= bits; // Decrement remaining bits to encode by the number encoded here
			
			// Same for G component of the RGB
			int g = rgb >> 8;
			if (bits_to_encode > 0)
			{
				g &= 0xFF;
				int data_g = (data << (32 - bits)) >>> (32 - bits);
			//	System.out.println("Wants to encode "+data_g + " into "+g);
				g >>>= (bits); g <<= bits; // Set lower 'bits' bits of value to 0 so we can do an or
				g |= data_g; // Encode data_g bits.  Hopefully, only "bits" bits exist.
			//	System.out.println("Comes out with "+g);
				data >>>= bits;
				bits_to_encode -= bits;
			}
			
			// Same for B component of the RGB
			int b = rgb & 0xFF;
			if (bits_to_encode > 0)
			{
				int data_b = (data << (32 - bits)) >>> (32 - bits);
			//	System.out.println("Wants to encode "+data_b + " into "+b);
				b >>>= (bits); b <<= bits; // Set lower 'bits' bits of value to 0 so we can do an or
				b |= data_b; // Encode data_b bits.  Hopefully, only "bits" bits exist.
			//	System.out.println("Comes out with "+b);
				data >>>= bits;
				bits_to_encode -= bits;
			}
			
			
			// Consolidate r, g, and b back into rgb
			
			rgb &=0xFF000000; // Cut off lower bits so we can replace them 
			rgb |= (r << 16);
			rgb |= (g << 8);
			rgb |= (b << 0);
			
			//System.out.println("Aftermath:"+Integer.toString(rgb,16));
			//System.out.printf("Changing offset=%d (%d, %d) to value=%X\n",offset -1, p.x, p.y, rgb);
			image.setRGB(p.x, p.y, rgb);
			
		
		}
		
		return offset;
		
	}
	
	Point linearToPoint(long offset)
	{
		long y = offset / image.getWidth();
		long x = offset - (image.getWidth()*y);
		
		return new Point((int)x, (int)y);
	}
	
	
}
